# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
# This file is part of ByteQC.
#
# Licensed under the Apache License, Version 2.0 (the "License")
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https: // www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pyscf import gto
from byteqc.cump2 import dfmp2_grad
from gpu4pyscf import scf as gscf


if __name__ == '__main__':
    mol = gto.M()
    mol.atom = '''
        O	-1.37129	1.46014	-1.47602
        O	-1.46014	-1.37129	1.47602
        O	1.46014	1.37129	1.47602
        O	-1.29921	-1.38806	-1.36527
        O	-1.38806	1.29921	1.36527
        O	1.29921	1.38806	-1.36527
        O	1.37129	-1.46014	-1.47602
        O	1.38806	-1.29921	1.36527
        H	-1.71999	2.10566	-2.10269
        H	-0.37943	1.52902	-1.53239
        H	-2.10566	-1.71999	2.10269
        H	-1.52902	-0.37943	1.53239
        H	1.52902	0.37943	1.53239
        H	2.10566	1.71999	2.10269
        H	-1.50726	-1.52116	-0.41981
        H	-1.50884	-0.44611	-1.51765
        H	-0.44611	1.50884	1.51765
        H	-1.52116	1.50726	0.41981
        H	1.50726	1.52116	-0.41981
        H	1.50884	0.44611	-1.51765
        H	1.71999	-2.10566	-2.10269
        H	0.37943	-1.52902	-1.53239
        H	1.52116	-1.50726	0.41981
        H	0.44611	-1.50884	1.51765
    '''
    mol.basis = 'ccpvdz'
    mol.verbose = 5
    mol.build()

    mf = gscf.RHF(mol)
    mf.kernel()
    e_corr, e_tot, grad = dfmp2_grad.kernel(
        mol, mf, auxbasis=f'{mol.basis}-ri')

# cc-pVDZ
# HF -608.328626525232
# MP2 correlation energy: -1.6900054024845605

# Note that the gradient results will be effect by the mf methods selected.
# the mf generated by gpu4pyscf.scf.RHF and pyscf.scf.RHF cause slight different grad results.
# the below grad is generated by pyscf.scf.RHF
#    0  O     0.0003299722     0.0028016886    -0.0030501476
#    1  O    -0.0028016886     0.0003299722     0.0030501476
#    2  O     0.0028016886    -0.0003299722     0.0030501476
#    3  O    -0.0004614286     0.0011517363     0.0012068276
#    4  O     0.0011517363     0.0004614286    -0.0012068276
#    5  O     0.0004614286    -0.0011517363     0.0012068276
#    6  O    -0.0003299722    -0.0028016886    -0.0030501476
#    7  O    -0.0011517363    -0.0004614286    -0.0012068276
#    8  H    -0.0022206453    -0.0000750937     0.0000737608
#    9  H     0.0006132940    -0.0002948647     0.0002310221
#   10  H     0.0000750937    -0.0022206453    -0.0000737608
#   11  H     0.0002948647     0.0006132940    -0.0002310221
#   12  H    -0.0002948647    -0.0006132940    -0.0002310221
#   13  H    -0.0000750937     0.0022206453    -0.0000737608
#   14  H    -0.0001002904    -0.0021800192     0.0003584740
#   15  H    -0.0002362485     0.0003622348    -0.0017192868
#   16  H     0.0003622348     0.0002362485     0.0017192868
#   17  H    -0.0021800192     0.0001002904    -0.0003584740
#   18  H     0.0001002904     0.0021800192     0.0003584740
#   19  H     0.0002362485    -0.0003622348    -0.0017192868
#   20  H     0.0022206453     0.0000750937     0.0000737608
#   21  H    -0.0006132940     0.0002948647     0.0002310221
#   22  H     0.0021800192    -0.0001002904    -0.0003584740
#   23  H    -0.0003622348    -0.0002362485     0.0017192868

# CPU results, slightly different from GPU results,
# becuase the int2e in gpu4pyscf is little bit different from pyscf.
#    0  O    3.29967220e-04  2.80168084e-03 -3.05014035e-03
#    1  O   -2.80168084e-03  3.29967220e-04  3.05014035e-03
#    2  O    2.80168084e-03 -3.29967220e-04  3.05014035e-03
#    3  O   -4.61429362e-04  1.15174006e-03  1.20682527e-03
#    4  O    1.15174006e-03  4.61429362e-04 -1.20682527e-03
#    5  O    4.61429362e-04 -1.15174006e-03  1.20682527e-03
#    6  O   -3.29967220e-04 -2.80168084e-03 -3.05014035e-03
#    7  O   -1.15174006e-03 -4.61429362e-04 -1.20682527e-03
#    8  H   -2.22064770e-03 -7.50855035e-05  7.37528894e-05
#    9  H    6.13302658e-04 -2.94866056e-04  2.31023013e-04
#   10  H    7.50855035e-05 -2.22064770e-03 -7.37528894e-05
#   11  H    2.94866056e-04  6.13302658e-04 -2.31023013e-04
#   12  H   -2.94866056e-04 -6.13302658e-04 -2.31023013e-04
#   13  H   -7.50855035e-05  2.22064770e-03 -7.37528894e-05
#   14  H   -1.00289978e-04 -2.18002021e-03  3.58476892e-04
#   15  H   -2.36246661e-04  3.62232736e-04 -1.71928636e-03
#   16  H    3.62232736e-04  2.36246661e-04  1.71928636e-03
#   17  H   -2.18002021e-03  1.00289978e-04 -3.58476892e-04
#   18  H    1.00289978e-04  2.18002021e-03  3.58476892e-04
#   19  H    2.36246661e-04 -3.62232736e-04 -1.71928636e-03
#   20  H    2.22064770e-03  7.50855035e-05  7.37528894e-05
#   21  H   -6.13302658e-04  2.94866056e-04  2.31023013e-04
#   22  H    2.18002021e-03 -1.00289978e-04 -3.58476892e-04
#   23  H   -3.62232736e-04 -2.36246661e-04  1.71928636e-03

# obtained by pyscf.grad.mp2 which is using MP2 without denisty fitting
# --------------- MP2 gradients ---------------
#          x                y                z
# 0 O     0.0003291478     0.0027915853    -0.0030412578
# 1 O    -0.0027915853     0.0003291478     0.0030412578
# 2 O     0.0027915853    -0.0003291478     0.0030412578
# 3 O    -0.0004650578     0.0011406070     0.0011959280
# 4 O     0.0011406070     0.0004650578    -0.0011959280
# 5 O     0.0004650578    -0.0011406070     0.0011959280
# 6 O    -0.0003291478    -0.0027915853    -0.0030412578
# 7 O    -0.0011406070    -0.0004650578    -0.0011959280
# 8 H    -0.0022029292    -0.0000747218     0.0000739996
# 9 H     0.0005907805    -0.0002865347     0.0002229040
# 10 H     0.0000747218    -0.0022029292    -0.0000739996
# 11 H     0.0002865347     0.0005907805    -0.0002229040
# 12 H    -0.0002865347    -0.0005907805    -0.0002229040
# 13 H    -0.0000747218     0.0022029292    -0.0000739996
# 14 H    -0.0001019540    -0.0021614765     0.0003497033
# 15 H    -0.0002371534     0.0003535661    -0.0017014043
# 16 H     0.0003535661     0.0002371534     0.0017014043
# 17 H    -0.0021614765     0.0001019540    -0.0003497033
# 18 H     0.0001019540     0.0021614765     0.0003497033
# 19 H     0.0002371534    -0.0003535661    -0.0017014043
# 20 H     0.0022029292     0.0000747218     0.0000739996
# 21 H    -0.0005907805     0.0002865347     0.0002229040
# 22 H     0.0021614765    -0.0001019540    -0.0003497033
# 23 H    -0.0003535661    -0.0002371534     0.0017014043
# ----------------------------------------------
